from flask import Flask, request, jsonify
from flask_cors import CORS
from queue import Queue
from threading import Thread
import datetime
import os
import pymysql
import requests
import json
from openai import OpenAI
import os
import json
import re

app = Flask(__name__)
CORS(app)

# MySQL 连接
db = pymysql.connect(
    host="localhost",
    user="root",
    password="0520Hlw!@#",
    database="aisql",
    charset="utf8mb4"
)
cursor = db.cursor()

# 队列
task_queue = Queue()

# 保存图像
def save_raw_image(upload_time, image_base64):
    sql = "INSERT INTO raw_images (upload_time, image_base64) VALUES (%s, %s)"
    cursor.execute(sql, (upload_time, image_base64))
    db.commit()
    return cursor.lastrowid

# 保存设备信息
def save_device_info(image_id, device_id, device_style):
    sql = "INSERT INTO device_info (image_id, device_id, device_style) VALUES (%s, %s, %s)"
    cursor.execute(sql, (image_id, device_id, device_style))
    db.commit()



# 初始化客户端
client = OpenAI(
    api_key="sk-728ab8263bf64d5d925bc3179148549f",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

def call_qwen_vl(base64_image):
    try:
        completion = client.chat.completions.create(
            model="qwen-vl-max",
            messages=[
                {
                    "role": "system",
                    "content": [{"type": "text", "text": "你是一个图像识别专家，负责提取设备信息。"}]
                },
                {
                    "role": "user",
                    "content": [
                        {"type": "image_url", "image_url": {"url": base64_image}},
                        {"type": "text", "text": "请提取图像中的功能位置和缺陷描述，并以 JSON 数组格式返回，例如：[{'功能位置': 'xxx', '缺陷描述': 'yyy'}]，不要输出 Markdown 代码块"}
                    ]
                }
            ]
        )

        content = completion.choices[0].message.content.strip()
        print(f"[模型返回原始内容] {content}")


        # ✅ 去掉 Markdown 代码块
        if content.startswith("```json") or content.startswith("```"):
            content = re.sub(r"^```(?:json)?\s*", "", content)
            content = re.sub(r"\s*```$", "", content)

        print("[清理后的 JSON]", content)

        # ✅ 提取并解析 JSON 数组
        data_array = json.loads(content)

        print("[模型识别结果类型]", type(data_array))

        results = []
        if isinstance(data_array, list):
            for item in data_array:
                results.append({
                    "device_id": item.get("功能位置", "未知"),
                    "device_style": item.get("缺陷描述", "未知")
                })
        elif isinstance(data_array, dict):
            # 万一是单个对象
            results.append({
                "device_id": data_array.get("功能位置", "未知"),
                "device_style": data_array.get("缺陷描述", "未知")
            })
        else:
            print("[解析失败] 返回值既不是 list 也不是 dict")

        return results

    except Exception as e:
        print(f"[模型调用异常] {e}")
        return [{
            "device_id": "识别失败",
            "device_style": "识别失败"
        }]


# # 模型 API 调用
# def call_qwen_vl(base64_image):
#     headers = {
#         "Authorization": f"Bearer {os.getenv('QWEN_API_KEY')}"
#     }
#     payload = {
#         "model": "qwen-vl-max-latest",
#         "messages": [
#             {
#                 "role": "user",
#                 "content": [
#                     {"type": "image_url", "image_url": {"url": base64_image}},
#                     {"type": "text", "text": "请提取图像中的功能位置和缺陷描述，并以 JSON 格式返回，不要输出 Markdown 代码块"}
#                 ]
#             }
#         ]
#     }
#     response = requests.post(
#         "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation",
#         headers=headers,
#         json=payload,
#         timeout=30
#     )
#     return response.json()

# 后台线程
def worker():
    while True:
        image_id, base64_image = task_queue.get()
        try:
            response = call_qwen_vl(base64_image)
            # 遍历返回的识别结果数组
            for item in response:
                device_id = item.get("device_id", "未知位置")
                device_style = item.get("device_style", "未知样式")
                save_device_info(image_id, device_id, device_style)
        except Exception as e:
            print(f"模型处理异常: {e}")
        task_queue.task_done()


Thread(target=worker, daemon=True).start()

@app.route("/api/upload", methods=["POST"])
def upload():
    data = request.get_json()
    image_base64 = data.get("image_base64")
    upload_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    image_id = save_raw_image(upload_time, image_base64)
    task_queue.put((image_id, image_base64))
    return jsonify({"message": "图像已接收", "image_id": image_id})

if __name__ == "__main__":
    app.run(port=5000)
